// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/testing/exec/util/sum_non_pod_aggregate.h>
#include <pollux/exec/aggregate.h>
#include <pollux/exec/hash_aggregation.h>
#include <pollux/expression/function_signature.h>

namespace kumo::pollux::exec::test {

int NonPODInt64::constructed = 0;
int NonPODInt64::destructed = 0;

namespace {

// SumNonPODAggregate uses NonPODInt64 as accumulator which has external memory
// NonPODInt64::constructed and NonPODInt64::destructed. By asserting their
// equality, we make sure Pollux calls constructor/destructor properly.
class SumNonPODAggregate : public Aggregate {
 public:
  explicit SumNonPODAggregate(pollux::TypePtr resultType, int alignment)
      : Aggregate(resultType), alignment_(alignment) {}

  int32_t accumulatorAlignmentSize() const override {
    return alignment_;
  }

  int32_t accumulatorFixedWidthSize() const override {
    return sizeof(NonPODInt64);
  }

  bool accumulatorUsesExternalMemory() const override {
    return true;
  }

  void extractAccumulators(
      char** groups,
      int32_t numGroups,
      pollux::VectorPtr* result) override {
    auto vector = (*result)->as<FlatVector<int64_t>>();
    vector->resize(numGroups);
    int64_t* rawValues = vector->mutableRawValues();
    uint64_t* rawNulls = getRawNulls(vector);
    for (int32_t i = 0; i < numGroups; ++i) {
      char* group = groups[i];
      if (isNull(group)) {
        vector->set_null(i, true);
      } else {
        clearNull(rawNulls, i);
        rawValues[i] = value<NonPODInt64>(group)->value;
      }
    }
  }

  void extractValues(char** groups, int32_t numGroups, pollux::VectorPtr* result)
      override {
    extractAccumulators(groups, numGroups, result);
  }

  void addIntermediateResults(
      char** groups,
      const pollux::SelectivityVector& rows,
      const std::vector<pollux::VectorPtr>& args,
      bool /*mayPushdown*/) override {
    DecodedVector decoded(*args[0], rows);

    rows.applyToSelected([&](vector_size_t i) {
      if (decoded.is_null_at(i)) {
        return;
      }
      clearNull(groups[i]);
      value<NonPODInt64>(groups[i])->value += decoded.value_at<int64_t>(i);
    });
  }

  void addRawInput(
      char** groups,
      const pollux::SelectivityVector& rows,
      const std::vector<pollux::VectorPtr>& args,
      bool mayPushdown) override {
    addIntermediateResults(groups, rows, args, mayPushdown);
  }

  void addSingleGroupIntermediateResults(
      char* group,
      const pollux::SelectivityVector& rows,
      const std::vector<pollux::VectorPtr>& args,
      bool /*mayPushdown*/) override {
    DecodedVector decoded(*args[0], rows);

    rows.applyToSelected([&](vector_size_t i) {
      if (decoded.is_null_at(i)) {
        return;
      }
      clearNull(group);
      value<NonPODInt64>(group)->value += decoded.value_at<int64_t>(i);
    });
  }

  void addSingleGroupRawInput(
      char* group,
      const pollux::SelectivityVector& rows,
      const std::vector<pollux::VectorPtr>& args,
      bool mayPushdown) override {
    addSingleGroupIntermediateResults(group, rows, args, mayPushdown);
  }

 protected:
  void initializeNewGroupsInternal(
      char** groups,
      melon::Range<const pollux::vector_size_t*> indices) override {
    for (auto i : indices) {
      char* group = value<char>(groups[i]);
      POLLUX_CHECK_EQ(reinterpret_cast<uintptr_t>(group) % alignment_, 0);
      new (group) NonPODInt64(0);
    }
  }

  void destroyInternal(melon::Range<char**> groups) override {
    for (auto group : groups) {
      if (isInitialized(group)) {
        value<NonPODInt64>(group)->~NonPODInt64();
      }
    }
  }

 private:
  const int32_t alignment_;
};

} // namespace

exec::AggregateRegistrationResult registerSumNonPODAggregate(
    const std::string& name,
    int alignment) {
  std::vector<std::shared_ptr<pollux::exec::AggregateFunctionSignature>>
      signatures{
          pollux::exec::AggregateFunctionSignatureBuilder()
              .returnType("bigint")
              .intermediateType("bigint")
              .argumentType("bigint")
              .build(),
      };

  return pollux::exec::registerAggregateFunction(
      name,
      std::move(signatures),
      [alignment](
          pollux::core::AggregationNode::Step /*step*/,
          const std::vector<pollux::TypePtr>& /*argTypes*/,
          const pollux::TypePtr& /*resultType*/,
          const core::QueryConfig& /*config*/)
          -> std::unique_ptr<pollux::exec::Aggregate> {
        return std::make_unique<SumNonPODAggregate>(pollux::BIGINT(), alignment);
      },
      false /*registerCompanionFunctions*/,
      true /*overwrite*/);
}

} // namespace kumo::pollux::exec::test
